# 处理数据集
dsTrain = pd.read_csv(r'mnist.csv')
# 原本kaggle上的测试集是没有lable分类的，在kaggle的leaderBoard上找到正确的lable后pd.concat到原本的测试集
dsTest = pd.read_csv(r'mnist_test.csv')
'''
# 获取训练集的特征
ds_std = np.std(dsTrain.iloc[:, 1:], axis = 0) # 训练集标准差
ds_mean = np.mean(dsTrain.iloc[:, 1:], axis = 0) # 训练集均值
# 标准化数据，此方法是针对训练集中的同一位置的像素，效果一般，故未采用
for i in range(len(ds.T) - 1):
    if ds_std[i] != 0: # 边角上就很可能出现标准差为0的像素
        dsR.iloc[:, i + 1] = (dsR.iloc[:, i + 1] - ds_mean[i]) / ds_std[i] # 标准分数公式，基本呈正态分布
'''
trainX = dsTrain.iloc[:, 1:]
trainy = dsTrain.iloc[:, :1]
# 排列测试集，方便在测试之后绘图
testX = pd.concat([dsTest[dsTest["label"] == i].iloc[:, 1:] for i in range(10)])
testy = pd.concat([dsTest[dsTest["label"] == i].iloc[:, :1] for i in range(10)])
# 归一化
trainX = trainX / 255
testX = testX / 255
# 获取训练集特征
mean = np.mean(np.array(trainX).ravel())  # 训练集均值
std = np.std(np.array(trainX).ravel())  # 训练集标准差
trainX = (trainX - mean) / std  # 标准分数公式，基本呈正态分布
testX = (testX - mean) / std  # 用训练集的特征处理测试集数据，才能证明模型是否能识别陌生数据

# d_func 是func的导数/偏导 derivative
def relu(x):
    return np.maximum(x, 0)

def d_relu(x):
    return x > 0

def lrelu(x):
    # leaky relu
    return np.maximum(x, 0.1 * x)

def d_lrelu(x):
    return (x > 0) * 0.9 + 0.1

def softMax(x):
    return np.array(np.exp(x) / np.sum(np.exp(x), axis=1))

def d_softMax(x):
    x = softMax(x)
    r = []
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            r.append([])
            for k in range(x.shape[1]):
                r[-1].append(-x[i][j] * x[i][k] if j != k else x[i][j] * (1 - x[i][k]))
    return np.array(r)

def sigmoid(x):
    x_ravel = np.array(x).ravel()
    length = len(x_ravel)
    y = []
    for i in range(length):
        if x_ravel[i] >= 0:
            x_ravel[i] = min(19, x_ravel[i])  # 限制范围，精度不够，np.exp()也易爆
            y.append(1.0 / (1 + np.exp(-x_ravel[i])))
        else:
            x_ravel[i] = max(-744, x_ravel[i])
            y.append(np.exp(x_ravel[i]) / (np.exp(x_ravel[i]) + 1))
    return np.array(y).reshape(x.shape)

def d_sigmoid(x):
    return sigmoid(x) * (1 - sigmoid(x))

def OneHot(length, width, y):
    # OneHot encoding
    y = np.array(y).ravel()
    r = np.zeros((length, width))
    for i in range(length):
        r[i][int(y[i])] = 1.0
    return r

def AddBias(x):
    # 在第0列插入全为1的列
    x = np.matrix(x)
    return np.c_[np.ones((x.shape[0], 1)), x]

def RemoveBias(x):
    # 去掉第一列
    return np.matrix(x)[:, 1:]

class FCNN:
    def Init(self, unitNum, dropOutProb=[]):
        '''
        完成一些初始化
        unitNum: list 每一层的神经元数量
        dropOutProb: list 除了输出层以外每一层的dropOut概率(0,1)
        '''
        self.unitNum = unitNum
        self.layerNum = len(unitNum)
        self.dropOutProb = dropOutProb if dropOutProb else [0 for i in range(self.layerNum - 1)]
        self.outLength = unitNum[-1]
        self.w = [0]  # 0 的作用仅为占位
        self.Z = [0]
        self.A = [0]
        self.cost = []
        self.costTemp = []
        self.dwNorm = []
        self.dwNormTemp = []
        # He initialization，针对隐藏层是reLu激活函数的权重初始化方法
        # 正确的权重初始化方法能一定程度避免梯度消失
        for i in range(1, self.layerNum - 1):
            np.random.seed(i)  # 设置随机数种子
            self.w.append(np.matrix(np.random.randn(unitNum[i - 1] + 1, unitNum[i])) * np.sqrt(2 / unitNum[i]))
            self.w[-1][0] = np.matrix(np.zeros((1, unitNum[i]))) # bias偏置初始化为0
        # 最后的权重初始化和之前的不太一样，需要更加接近0，这是针对sigmoid/softMax
        np.random.seed(self.layerNum - 1)
        self.w.append(np.matrix(np.random.rand(unitNum[-2] + 1, unitNum[-1])))
        self.w[-1] = (self.w[-1] - np.mean(self.w[-1])) / np.std(self.w[-1]) / 100
        self.w[-1][0] = np.matrix(np.zeros((1, unitNum[-1])))
        self.ClrTempResult()

    def PrepareBatch(self, X, y, batchSize=0):
        '''
        分割训练集并完成一些初始化
        batchSize: 每一批样本的数量，需能被样本总数整除
        '''
        N = X.shape[0]
        if not batchSize:
            batchSize = N
        if N % batchSize:
            print("# ERR: illegal batchSize")
            return
        self.batchSize = batchSize
        self.batchTime = int(N / batchSize)
        self.X = np.matrix(X)
        self.y = OneHot(X.shape[0], unitNum[-1], y)
        self.batchX = []
        self.batchy = []
        self.batchTime = int(N / batchSize)
        for i in range(self.batchTime):
            # mini-batch可以加快训练速度
            self.batchX.append(np.array(self.X[i * batchSize : (i + 1) * batchSize]))
            self.batchy.append(self.y[i * batchSize : (i + 1) * batchSize])
            print("\rPreparing: %.2f%%" % ((i + 1) / N * batchSize * 100), end="")

    def fit(self, epoch, rate):
        '''
        epoch: int 训练轮次
        rate: float 学习率
        '''
        N = self.X.shape[0]
        for i in range(epoch):
            for j in range(self.batchTime):
                self.Z[0] = self.batchX[j]
                self.A[0] = self.batchX[j]
                self.y = self.batchy[j]
                self.gradDes(rate)
                self.ClrTempResult()
                print(
                    "\rTotalProgress: %.2f%% BatchProgress: %.2f%%  "
                    % (((i + 1) / epoch * 100), ((j + 1) / N * self.batchSize * 100)),
                    end="",
                )
            # 记录一个轮次下的平均损失函数
            self.cost.append(sum(self.costTemp) / self.batchTime)
            self.costTemp.clear()
            # 记录一个轮次下的平均dw的L2，即模长
            self.dwNorm.append(np.sum(self.dwNormTemp, axis=0) / self.batchTime)
            self.dwNormTemp.clear()

    def gradDes(self, rate):
        self.backProp()
        # 没有加入正则项惩罚，因为未发现过拟合现象，同时样本足够多时可以不加入正则项
        '''
        # 一种较为简单的损失函数 J(θ)
        J = np.sum(np.multiply(self.y-self.A[-1],self.y-self.A[-1]))/self.batchSize
        # 这里使用更加复杂的带有log的交叉熵函数
        '''
        J = (-np.sum(np.array(self.y) * np.log(self.A[-1]) + np.array(1 - self.y) * np.log(1 - self.A[-1])) / self.batchSize)
        self.costTemp.append(J)
        dwNt = []
        for i in range(1, len(self.w)):
            dwNt.append(np.linalg.norm(self.dw[i]))
            self.w[i] -= rate * self.dw[i]
        self.dwNormTemp.append(dwNt)

    def GetDropOut(self):
        # 获取dropOut矩阵，根据给定概率随机去掉神经元
        index = []
        prob = self.dropOutProb[0]
        num = 0 if not prob else int(np.clip(np.random.randn() / 50 + prob, 0, 0.75) * self.unitNum[0])
        index.append(np.random.choice(self.unitNum[0], num, False))
        for i in range(1, self.layerNum - 1):
            prob = self.dropOutProb[i]
            num = 0 if not prob else int(np.clip(np.random.randn() / 50 + prob, 0, 0.75) * self.unitNum[i])
            index.append(np.random.choice(self.unitNum[i], num, False))
            dropOutMat = np.ones((self.unitNum[i - 1], self.unitNum[i]))
            for j in index[i - 1]:
                dropOutMat[j] = np.zeros((1, self.unitNum[i]))
            dropOutMat = np.transpose(dropOutMat)
            for j in index[i]:
                dropOutMat[j] = np.zeros((1, self.unitNum[i - 1]))
            dropOutMat = np.r_[np.ones((1, self.unitNum[i])), np.transpose(dropOutMat)]
            self.dropOut.append(np.matrix(dropOutMat))
        dropOutMat = np.ones((self.unitNum[-2], self.unitNum[-1]))
        for j in index[-1]:
            dropOutMat[j] = np.zeros((1, self.unitNum[-1]))
        dropOutMat = np.r_[np.ones((1, self.unitNum[-1])), dropOutMat]
        self.dropOut.append(np.matrix(dropOutMat))

    def backProp(self, d_lossFunc, forward=True, dX=False):
        # 有详细推导过程
        if forward:
            self.forwardProp()
        self.dA[-1] = d_lossFunc()
        # --- appaly derivative of SoftMax
        self.dZ[-1] = []
        Zt = np.array(self.Z[-1])
        dAt = np.array(self.dA[-1])
        for i in range(self.batchSize):
            self.dZ[-1].append(np.array(d_softMax(np.matrix(Zt[i])) *
                                   np.matrix(dAt[i]).reshape((self.unitNum[-1], 1))).ravel())
        self.dZ[-1] = np.matrix(self.dZ[-1])
        # ---
        dw = self.A[-2].T * self.dZ[-1]
        db = np.sum(self.dZ[-1], axis=0)
        self.dw[-1] = np.r_[db, dw] / self.batchSize
        for i in reversed(range(1, self.layerNum - 1)):
            self.dA[i] = RemoveBias(self.dZ[i + 1] * np.multiply(self.w[i + 1], self.dropOut[i]).T)
            self.dZ[i] = np.multiply(d_relu(self.Z[i]), self.dA[i])
            dw = self.A[i - 1].T * self.dZ[i]
            db = np.sum(self.dZ[i], axis=0)
            self.dw[i] = np.multiply(np.r_[db, dw] / self.batchSize, self.dropOut[i - 1])
        if dX:
            self.dA[0] = RemoveBias(self.dZ[1] * np.multiply(self.w[1], self.dropOut[0]).T)
            self.dZ[0] = self.dA[0]

    def forwardProp(self):
        self.GetDropOut()
        for i in range(1, self.layerNum - 1):
            self.Z.append(AddBias(self.A[-1]) * np.multiply(self.w[i], self.dropOut[i - 1]))
            self.A.append(relu(self.Z[-1]))
        self.Z.append(AddBias(self.A[-1]) * np.multiply(self.w[-1], self.dropOut[-1]))
        self.A.append(softMax(self.Z[-1]))

    def ClrTempResult(self):
        self.Z = self.Z[0:1]
        self.A = self.A[0:1]
        self.dw = list(np.repeat(1, self.layerNum, axis=0))
        self.dZ = list(np.repeat(1, self.layerNum, axis=0))
        self.dA = list(np.repeat(1, self.layerNum, axis=0))
        self.dropOut = []

    def predict(self, X, y):
        self.ClrTempResult()
        self.dropOutProb = [0 for i in range(self.layerNum - 1)]
        self.Z[0] = np.matrix(X)
        self.A[0] = np.matrix(X)
        self.forwardProp()
        r_p = self.A[-1]
        self.r = np.zeros((r_p.shape[0], 1))
        correct = 0
        self.wrong = []
        for i in range(r_p.shape[0]):
            self.r[i] = np.argmax(r_p[i])
            if self.r[i] == np.array(y).ravel()[i]:
                correct += 1
            else:
                self.wrong.append(i) # 储存错误样本的index
        print("correct: " + str(correct / r_p.shape[0]))

    def SaveParameters(self):
        f = open("FCNNweight.dat", "wb")
        wt = []
        for i in self.w[1:]:
            wt.extend(list(np.array(i).ravel()))
        f.write(np.array(wt).tobytes())
        f.close()

    def LoadParameters(self):
        f = open("FCNNweight.dat", "rb")
        wt = np.frombuffer(f.read(), np.float64)
        f.close()
        for i in range(1, len(self.unitNum)):
            self.w[i] = np.matrix(wt[: (self.unitNum[i - 1] + 1) * self.unitNum[i]]).reshape((self.unitNum[i - 1] + 1, self.unitNum[i]))
            wt = wt[(self.unitNum[i - 1] + 1) * self.unitNum[i] :]

def ShowCostFig(epoch, cost):
    # 绘制损失函数在迭代中的变化图像
    ax = plt.figure().gca()
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))  # 使坐标刻度为整数
    plt.plot(range(0, epoch), cost[0:], "go-")
    plt.show()
    print("Final Cost: " + str(cost[epoch - 1]))

def ShowdwFig(epoch, dw, layerNum):
    # 绘制dw的L2在迭代中的变化图像
    camp = getCmap(layerNum)
    ax = plt.figure(figsize=(15, 5)).gca()
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    for i, data in enumerate(np.transpose(dw)):
        plt.plot(data[1:], c=camp(i), label=i)
    plt.legend()  # 加上图例
    plt.show()

def getCmap(n, name="hsv"):
    # 从网上copy了一个调色盘来选颜色
    """Returns a function that maps each index in 0, 1, ..., n-1 to a distinct
    RGB color; the keyword argument name must be a standard mpl colormap name."""
    return plt.cm.get_cmap(name, n)

def ShowNorm(x):
    print("Norm: " + str(np.linalg.norm(x)))